import argparse
from dataclasses import dataclass
import logging
import subprocess
import time
from typing import Tuple

from kubernetes import config, client
import psutil as ps
import yaml

from auto_redis_client import RedisClientParams, run_redis_client_benchmark
from util import check_resource_num

config.load_kube_config()
core_v1 = client.CoreV1Api()
apps_v1 = client.AppsV1Api()

logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)s: %(message)s",
)


def cleanup_server(deploy, svc: dict):
    deploy_name, svc_name = deploy["metadata"]["name"], svc["metadata"]["name"]
    deploy_ns, svc_ns = deploy["metadata"]["namespace"], svc["metadata"]["namespace"]
    apps_v1.delete_namespaced_deployment(deploy_name, deploy_ns)
    core_v1.delete_namespaced_service(svc_name, svc_ns)
    logging.info(f"已清除 Deployment {deploy_name} & Service {svc_name}")


def read_server_obj(template_job: str) -> Tuple[dict, dict]:
    deploy, svc = {}, {}
    with open(template_job, "r") as f:
        for obj in yaml.safe_load_all(f):
            if obj["kind"] == "Deployment":
                deploy = obj
            elif obj["kind"] == "Service":
                svc = obj
            else:
                pass
    return deploy, svc


def update_deployment(
    body: dict,
    cpu_requests: str = "",
    cpu_limits: str = "",
) -> dict:
    if cpu_requests and not check_resource_num(cpu_requests):
        raise ValueError(f"cpu request is not valid: {cpu_requests}")
    if cpu_limits and not check_resource_num(cpu_limits):
        raise ValueError(f"cpu limits is not valid: {cpu_limits}")

    # template has empty resources by default, unless u modify it
    container = body["spec"]["template"]["spec"]["containers"][0]
    if cpu_requests or cpu_limits:
        if not container.get("resources", None):
            container["resources"] = {
                "requests": {"memory": "8Gi"},
                "limits": {"memory": "8Gi"},
            }
    if cpu_requests:
        container["resources"]["requests"]["cpu"] = cpu_requests
    if cpu_limits:
        container["resources"]["limits"]["cpu"] = cpu_limits

    return body


def install_redis_server(deploy: dict, svc: dict, wait_timeout: int = 60) -> bool:
    deploy_ns, svc_ns = deploy["metadata"]["namespace"], svc["metadata"]["namespace"]
    if deploy_ns != svc_ns:
        raise ValueError(f"Service 与 Deployment 命名空间需相同： {deploy_ns} != {svc_ns}")

    logging.info(f"开始部署 Redis server 资源:")
    logging.info(f"    Deployment: {deploy['metadata']['name']}")
    logging.info(f"    Service: {svc['metadata']['name']}")
    try:
        core_v1.create_namespaced_service(
            namespace=svc_ns,
            body=svc,
        )
    except Exception as e:
        logging.error(f"Service ${svc['metadata']['name']} 安装失败: {svc}")
        logging.error(f"\n{e}")
    try:
        apps_v1.create_namespaced_deployment(
            namespace=deploy_ns,
            body=deploy,
        )
    except Exception as e:
        logging.error(f"Deployment ${deploy['metadata']['name']} 安装失败: {deploy}")
        logging.error(f"{e.body}")

    try:
        time.sleep(1)  # wait for pod to be created
        cmd = [
            "kubectl", "wait",
            "pod",
            "--for=condition=ready",
            "-l", "app=redis-server",
            f"--namespace={deploy_ns}",
            f"--timeout={wait_timeout}s",
        ]
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError:
        return False

    # 获取pod列表并检查redis服务是否正常运行
    pods = core_v1.list_namespaced_pod(
        deploy_ns, label_selector="app=redis-server"
    )
    if not pods.items:
        return False
    pod_name = pods.items[0].metadata.name
    # 在pod中执行redis-cli ping命令
    redis_host = svc["metadata"]["name"]
    redis_port = str(svc["spec"]["ports"][0]["port"])
    cmd = [
        "kubectl", "exec",
        "-n", deploy_ns,
        pod_name,
        "--",
        "redis-cli",
        "-h", redis_host,
        "-p", redis_port,
        "ping",
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return result.stdout.strip() == "PONG"
    except subprocess.CalledProcessError:
        return False


def start_stress(stress_command: str, wait_time: int = 10) -> bool:
    stress_cmd = stress_command.strip().split(" ")
    try:
        subprocess.run(stress_cmd, check=True)
    except Exception as e:
        return False
    logging.info(f"等待加压生效，等待时间: {wait_time}s")
    time.sleep(wait_time)
    return True


def stop_stress(stress_command: str):
    if not stress_command:
        return
    stress_command = stress_command.strip()
    for proc in ps.process_iter(["pid", "name", "cmdline"]):
        cmdline = " ".join(proc.info["cmdline"]) if proc.info["cmdline"] else ""
        if stress_command in cmdline:
            proc.kill()
            logging.info(f"已结束加压进程 {proc.info['pid']}")


@dataclass
class RedisServerParams:
    run_time: int = 1
    test_type: str = "close"
    client_base_yaml: str = "redis/test-redis-client.yaml"
    server_base_yaml: str = "redis/test-redis-server.yaml"
    stress_command: str = ""
    server_cpu_requests: str = ""
    server_cpu_limits: str = ""
    client_cpu_requests: str = ""
    client_cpu_limits: str = ""
    redis_request_num: int = 1000000
    redis_client_num: int = 50
    redis_op_type: str = "get"
    redis_threads: int = 1


def run_redis_benchmark(params: RedisServerParams):
    logging.info("====开始 Redis 测试====")
    deploy, svc = read_server_obj(params.server_base_yaml)
    deploy = update_deployment(
        body=deploy,
        cpu_requests=params.server_cpu_requests,
        cpu_limits=params.server_cpu_limits,
    )
    if not install_redis_server(deploy, svc):
        logging.error("安装 Redis server 失败")
        return
    redis_host = svc["metadata"]["name"]
    redis_port = str(svc["spec"]["ports"][0]["port"])
    logging.info(f"Redis Service 运行于 {redis_host}:{redis_port}")

    if not params.stress_command:
        logging.info("未指定加压命令，跳过压力测试")
    else:
        logging.info(f"启动加压测试，加压命令：{params.stress_command}")
        if not start_stress(params.stress_command):
            logging.error("启动加压测试失败")
            stop_stress(params.stress_command)
            cleanup_server(deploy, svc)
            return

    try:
        client_params = RedisClientParams(
            run_time=params.run_time,
            test_type=params.test_type,
            base_yaml=params.client_base_yaml,
            cpu_requests=params.client_cpu_requests,
            cpu_limits=params.client_cpu_limits,
            redis_host=redis_host,
            redis_port=redis_port,
            redis_request_num=params.redis_request_num,
            redis_client_num=params.redis_client_num,
            redis_op_type=params.redis_op_type,
            redis_threads=params.redis_threads,
        )
        run_redis_client_benchmark(client_params)
    except Exception as e:
        logging.error("Redis benchmark 运行失败")
        stop_stress(params.stress_command)
        cleanup_server(deploy, svc)
        return

    stop_stress(params.stress_command)
    cleanup_server(deploy, svc)
    logging.info("====Redis 测试已完成====")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Redis服务端基准测试脚本")

    parser.add_argument(
        "-n",
        "--run-time",
        type=int,
        default=1,
        help="基准测试运行的次数",
    )
    parser.add_argument(
        "-t",
        "--test-type",
        type=str,
        default="close",
        help="测试类型(例如: close)",
    )
    parser.add_argument(
        "-sYaml",
        "--server-base-yaml",
        type=str,
        default="redis/test-redis-server.yaml",
        help="Redis服务端Deployment&Service的基础YAML文件",
    )
    parser.add_argument(
        "-cYaml",
        "--client-base-yaml",
        type=str,
        default="redis/test-redis-client.yaml",
        help="Redis客户端job的基础YAML文件",
    )
    parser.add_argument(
        "-stress-command",
        type=str,
        default="",
        help="加压测试命令，若为空则不进行加压测试",
    )
    parser.add_argument(
        "-sReq",
        "--server-cpu-requests",
        type=str,
        default="",
        help="服务端pod的CPU请求值",
    )
    parser.add_argument(
        "-sLimit",
        "--server-cpu-limits",
        type=str,
        default="",
        help="服务端pod的CPU限制值",
    )
    parser.add_argument(
        "-cReq",
        "--client-cpu-requests",
        type=str,
        default="",
        help="客户端pod的CPU请求值",
    )
    parser.add_argument(
        "-cLimit",
        "--client-cpu-limits",
        type=str,
        default="",
        help="客户端pod的CPU限制值",
    )
    parser.add_argument(
        "-r",
        "--redis-request-num",
        type=int,
        default=1000000,
        help="redis-benchmark的请求数量",
    )
    parser.add_argument(
        "-c",
        "--redis-client-num",
        type=int,
        default=50,
        help="redis-benchmark的客户端数量",
    )
    parser.add_argument(
        "--redis-op-type",
        type=str,
        default="get",
        help="Redis操作类型(例如: get, set)",
    )
    parser.add_argument(
        "--redis-threads",
        type=int,
        default=1,
        help="redis-benchmark的线程数量",
    )
    args = parser.parse_args()

    params = RedisServerParams(
        run_time=args.run_time,
        test_type=args.test_type,
        server_base_yaml=args.server_base_yaml,
        client_base_yaml=args.client_base_yaml,
        stress_command=args.stress_command,
        server_cpu_requests=args.server_cpu_requests,
        server_cpu_limits=args.server_cpu_limits,
        client_cpu_requests=args.client_cpu_requests,
        client_cpu_limits=args.client_cpu_limits,
        redis_request_num=args.redis_request_num,
        redis_client_num=args.redis_client_num,
        redis_op_type=args.redis_op_type,
        redis_threads=args.redis_threads,
    )
    run_redis_benchmark(params)
